package com.dman.multi.core;

import java.io.EOFException;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.util.Set;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.ReentrantLock;

public class ThrottledNetworkStream extends NetworkStream {
	
	private long prevTime;
	private long bytesRead;
	
	private static AtomicLong bytesPerSecond = new AtomicLong();	
	private static Mode mode = null;
	private static ReentrantLock rlock = new ReentrantLock(); 
	
	public static enum Mode{
		HEAVY(3), MEDIUM(2), LIGHT(1);
		
		private int number = 0;
		private static int maxNumber = 4;
		
		private Mode(int i) {
			number = i;
		}

		public int getNumber() {
			return number;
		}
		
		public static int getMaxNumber() {
			return maxNumber;
		}
	}
	
	public ThrottledNetworkStream(long bps) throws IOException {
		super();
		prevTime = System.nanoTime();
		rlock.lock();
		try{
			bytesPerSecond.set(bps);
			Mode md = mode;
			mode = Mode.HEAVY;
			if(md != null)
				setMode(md);
		} finally {
			rlock.unlock();			
		}
	}
	
	@Override
	public Set<SelectionKey> getReadySockets(long timeout) throws IOException, InterruptedException {
		Set<SelectionKey> setOfSock = null;

		long currTime = System.nanoTime();
		long theoreticalTime = (long) ((double) bytesRead / bytesPerSecond.get() * 1000000000);
		long actualTime = currTime - prevTime;
		if(actualTime < theoreticalTime - 1000000)
			Thread.sleep((theoreticalTime - actualTime)/1000000);

		setOfSock = super.getReadySockets(timeout);
		bytesRead = 0;
		prevTime = currTime;

		return setOfSock;
	}
	
	@Override
	public ByteBuffer readSocket(SelectionKey selKey) throws EOFException,
			IOException {
		ByteBuffer bb = super.readSocket(selKey);
		bytesRead += bb.limit() - bb.position();
		return bb;
	}

	public static void setMode(Mode md) {
		int modeNum = md.getNumber() + 1;
		int maxNum = Mode.getMaxNumber();
		rlock.lock();
		try{
			int prevModeNum = mode.getNumber() + 1;
			long connSpeed = bytesPerSecond.get() * maxNum / prevModeNum;
			bytesPerSecond.set(connSpeed * modeNum / maxNum);
			mode = md;	
		} finally {
			rlock.unlock();
		}
	}
}
